from matplotlib import pyplot as plot

import torch
from sklearn.model_selection import train_test_split

class Sample():
    def exec(self):
        self.prepare_data()
        self.train()
        self.predict()
        self.plot()
    
    def prepare_data(self):
        self.x = torch.linspace(1, 100, 100).type(torch.FloatTensor)
        rand = torch.randn(100) * 10
        self.y =  self.x + rand
        #self.data = train_test_split(self.x, self.y)
        self.data = self.x[:-10], self.x[-10:], self.y[:-10], self.y[-10:]
        self.a = torch.rand(1, requires_grad=True)
        self.b = torch.rand(1, requires_grad=True)
        self.learning_rate = 0.0001

    def train(self):
        for i in range(2000):
            prediction = self.a.expand_as(self.data[0]) * self.data[0] + self.b.expand_as(self.data[0])
            loss = torch.mean((prediction - self.data[2]) ** 2)
            if i % 200 == 0:
                print(f'loss: {loss}')
            loss.backward()
            self.a.data.add_( - self.learning_rate * self.a.grad.data)
            self.b.data.add_( - self.learning_rate * self.b.grad.data)
            self.a.grad.data.zero_()
            self.b.grad.data.zero_()
            
    def predict(self):
        self.pred = self.a.expand_as(self.data[1]) * self.data[1] + self.b.expand_as(self.data[1])
            
    def plot(self):
        plot.figure(figsize=(10, 8))
        plot.plot(self.data[0].data, self.data[2].data, 'o')
        plot.plot(self.data[1].data, self.data[3].data, 's')
        plot.plot(self.data[0], self.data[0] * self.a.data + self.b.data)
        plot.plot(self.data[1], self.pred.detach().numpy(), 'o')
        plot.xlabel('x')
        plot.ylabel('y')
        plot.show()

def main():
    Sample().exec()
    pass

if __name__ == '__main__':
    main()